package cn.bugstack.ai.test.spring.ai;

import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.client.McpSyncClient;
import io.modelcontextprotocol.client.transport.ServerParameters;
import io.modelcontextprotocol.client.transport.StdioClientTransport;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.mcp.SyncMcpToolCallbackProvider;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.test.context.SpringBootTest;

import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.List;

/**
 * 企业级智能研报生成器测试类
 *
 * 作用: 实现 Orchestrator-Workers 模式的深度研究Agent
 * 触发条件: 接收研究题目(如 "Spring AI vs LangChain 优劣势对比")
 * 数据流向: Manager → Worker1(搜索) → Worker2(文件) → Refiner → Reviewer → Finalizer
 */
@Slf4j
@SpringBootTest
public class DeepResearchAgentTest {

    /** DeepSeek API配置 */
    @Value("${spring.ai.openai.base-url}")
    private String openaiBaseUrl;

    @Value("${spring.ai.openai.api-key}")
    private String openaiApiKey;

    /** DeepSeek Chat 模型实例 */
    private ChatModel chatModel;

    /** MCP 工具回调提供器 */
    private SyncMcpToolCallbackProvider toolCallbackProvider;

    /** 输出目录(保存研报) */
    private static final String OUTPUT_DIR = "E:\\cursor\\ai-agent-station-study\\data\\reports";

    /**
     * 测试初始化(每个测试前执行)
     */
    @BeforeEach
    public void setUp() {
        log.info("\n========== 初始化 DeepResearch Agent 测试环境 ==========");

        // 初始化 DeepSeek Chat 模型(使用配置文件中的API密钥)
        OpenAiApi openAiApi = OpenAiApi.builder()
                .baseUrl(openaiBaseUrl)
                .apiKey(openaiApiKey)
                .completionsPath("v1/chat/completions")
                .embeddingsPath("v1/embeddings")
                .build();

        this.chatModel = OpenAiChatModel.builder()
                .openAiApi(openAiApi)
                .defaultOptions(OpenAiChatOptions.builder()
                        .model("deepseek-chat")
                        .maxTokens(8192)
                        .temperature(0.7)
                        .build())
                .build();

        log.info("✅ ChatModel 初始化完成 (Base URL: {}, API Key: {}...)", openaiBaseUrl, openaiApiKey.substring(0, 10));

        // 创建输出目录
        try {
            Files.createDirectories(Paths.get(OUTPUT_DIR));
            log.info("✅ 输出目录创建完成: {}", OUTPUT_DIR);
        } catch (IOException e) {
            log.error("❌ 输出目录创建失败", e);
        }
    }

    // ========== MCP 客户端工具方法 ==========

    /**
     * 创建文件系统 MCP 客户端
     *
     * @return McpSyncClient 实例
     */
    private McpSyncClient createFilesystemMcpClient() {
        log.info("🔧 创建文件系统 MCP 客户端...");

        String npxPath = System.getProperty("os.name").toLowerCase().contains("win")
                ? "npx.cmd" : "npx";

        ServerParameters stdioParams = ServerParameters.builder(npxPath)
                .args("-y", "@modelcontextprotocol/server-filesystem",
                      "E:\\\\cursor\\\\ai-agent-station-study")
                .build();

        McpSyncClient mcpClient = McpClient.sync(new StdioClientTransport(stdioParams))
                .requestTimeout(Duration.ofMinutes(3))
                .build();

        try {
            mcpClient.initialize();
            log.info("✅ 文件系统 MCP 客户端创建成功");
            return mcpClient;
        } catch (Exception e) {
            log.error("❌ 文件系统 MCP 客户端创建失败", e);
            throw new RuntimeException("Failed to create filesystem MCP client", e);
        }
    }

    /**
     * 创建百度AI搜索 MCP 客户端(带重试机制)
     *
     * @return McpSyncClient 实例
     */
    private McpSyncClient createBaiduAiSearchMcpClient() {
        log.info("🔧 创建百度AI搜索 MCP 客户端...");

        String apiKey = "bce-v3/ALTAK-HJfMMe4ChcyoaUAKhUJfk/0355338dc08c55b7c7bb1e5b88f3a0800c80bca1";
        String baseUrl = "http://appbuilder.baidu.com/v2/ai_search/mcp/";
        String sseEndpoint = "sse?api_key=" + apiKey;

        int maxRetries = 3;
        int retryDelayMs = 2000;
        Exception lastException = null;

        for (int attempt = 1; attempt <= maxRetries; attempt++) {
            try {
                log.info("尝试连接百度 AI 搜索 MCP (第 {}/{} 次)", attempt, maxRetries);

                io.modelcontextprotocol.client.transport.HttpClientSseClientTransport sseClientTransport =
                        io.modelcontextprotocol.client.transport.HttpClientSseClientTransport.builder(baseUrl)
                                .sseEndpoint(sseEndpoint)
                                .build();

                McpSyncClient mcpSyncClient = McpClient.sync(sseClientTransport)
                        .requestTimeout(Duration.ofMinutes(3))
                        .build();

                mcpSyncClient.initialize();
                log.info("✅ 百度AI搜索 MCP 客户端创建成功 (尝试 {}/{})", attempt, maxRetries);
                return mcpSyncClient;
            } catch (Exception e) {
                lastException = e;
                log.warn("⚠️ 百度AI搜索 MCP 客户端创建失败 (尝试 {}/{}): {}", attempt, maxRetries, e.getMessage());
                if (attempt < maxRetries) {
                    log.info("等待 {} 毫秒后重试...", retryDelayMs);
                    try {
                        Thread.sleep(retryDelayMs);
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                    }
                }
            }
        }

        log.error("❌ 百度AI搜索 MCP 客户端创建失败(已达最大重试次数)", lastException);
        throw new RuntimeException("Failed to create Baidu AI Search MCP client after retries", lastException);
    }

    // ========== 核心测试方法 ==========

    /**
     * 测试: 企业级智能研报生成器完整工作流
     *
     * 场景: 接收题目 "Spring AI vs LangChain 优劣势对比"
     * 流程: Manager → Worker1(搜索) → Worker2(文件) → Refiner → Reviewer → Finalizer
     */
    @Test
    public void test_deep_research_agent_workflow() {
        log.info("\n" +
                "╔═══════════════════════════════════════════════════════════════╗\n" +
                "║         企业级智能研报生成器 - 完整工作流测试                    ║\n" +
                "╚═══════════════════════════════════════════════════════════════╝");

        // ========== 阶段0: 定义研究题目 ==========
        String researchTopic = "Spring AI vs LangChain 优劣势对比";
        log.info("\n📋 研究题目: {}", researchTopic);
        log.info("⏰ 开始时间: {}\n", LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")));

        // ========== 阶段1: Manager Agent - 分析题目并拆解任务 ==========
        log.info("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
        log.info("🎯 阶段1: Manager Agent 分析题目");
        log.info("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");

        ChatClient managerAgent = ChatClient.builder(chatModel).build();

        String managerPrompt = String.format("""
                你是一个研究任务管理器(Manager Agent)。

                接收到的研究题目是: "%s"

                请分析这个题目，并拆解为2个子任务:
                1. 需要搜索哪些关键词来获取外部技术博客信息？(给出3-5个关键词)
                2. 需要从本地文件系统读取哪些相关文档？(给出文件路径或搜索模式)

                请用简洁的JSON格式返回:
                {
                  "search_keywords": ["关键词1", "关键词2", ...],
                  "local_files_patterns": ["*.md", "docs/**/*.pdf", ...]
                }
                """, researchTopic);

        String managerResponse = managerAgent.prompt()
                .user(managerPrompt)
                .call()
                .content();

        log.info("✅ Manager 分析结果:\n{}", managerResponse);

        // ========== 阶段2: Worker1 - 搜索外部信息 ==========
        log.info("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
        log.info("🔍 阶段2: Worker1 搜索外部技术博客");
        log.info("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");

        // 初始化搜索MCP客户端
        McpSyncClient searchClient = createBaiduAiSearchMcpClient();
        SyncMcpToolCallbackProvider searchTools = new SyncMcpToolCallbackProvider(searchClient);

        ChatClient worker1Agent = ChatClient.builder(chatModel)
                .defaultToolCallbacks(searchTools)
                .build();

        String worker1Prompt = String.format("""
                你是一个搜索专员(Worker1 Agent)。

                Manager 给你的任务是搜索关于 "%s" 的技术博客。
                Manager 的分析结果是:
                %s

                请使用搜索工具，针对每个关键词搜索相关技术博客(总共不超过10条结果)。
                将搜索结果整理为结构化列表:
                1. 博客标题
                2. 核心观点摘要(50字内)
                3. 来源链接

                注意: 只搜索真实存在的内容，不要编造。
                """, researchTopic, managerResponse);

        String worker1Response = worker1Agent.prompt()
                .user(worker1Prompt)
                .call()
                .content();

        log.info("✅ Worker1 搜索结果:\n{}", worker1Response);

        // ========== 阶段3: Worker2 - 读取本地文档 ==========
        log.info("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
        log.info("📂 阶段3: Worker2 读取本地文档");
        log.info("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");

        // 初始化文件系统MCP客户端
        McpSyncClient filesystemClient = createFilesystemMcpClient();
        SyncMcpToolCallbackProvider filesystemTools = new SyncMcpToolCallbackProvider(filesystemClient);

        ChatClient worker2Agent = ChatClient.builder(chatModel)
                .defaultToolCallbacks(filesystemTools)
                .build();

        String worker2Prompt = String.format("""
                你是一个文件系统专员(Worker2 Agent)。

                Manager 给你的任务是读取本地文档关于 "%s" 的内容。
                Manager 的分析结果是:
                %s

                请使用文件系统工具:
                1. 搜索项目中可能包含相关信息的文件(README.md, docs目录等)
                2. 读取这些文件的内容
                3. 提取关键信息并整理摘要

                如果找不到相关文件，请明确说明。
                """, researchTopic, managerResponse);

        String worker2Response = worker2Agent.prompt()
                .user(worker2Prompt)
                .call()
                .content();

        log.info("✅ Worker2 本地文档结果:\n{}", worker2Response);

        // ========== 阶段4: Refiner Agent - 汇总信息生成初稿 ==========
        log.info("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
        log.info("✍️ 阶段4: Refiner Agent 生成初稿");
        log.info("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");

        ChatClient refinerAgent = ChatClient.builder(chatModel).build();

        String refinerPrompt = String.format("""
                你是一个研报撰写专员(Refiner Agent)。

                研究题目: "%s"

                Worker1 的搜索结果:
                %s

                Worker2 的本地文档结果:
                %s

                请汇总以上信息，撰写一份完整的技术研报初稿(Markdown格式):

                # 标题
                ## 1. 背景介绍
                ## 2. 技术对比分析
                   - Spring AI 优势/劣势
                   - LangChain 优势/劣势
                ## 3. 使用场景建议
                ## 4. 参考资料

                要求:
                - 必须基于实际搜索到的内容撰写
                - 引用外部博客时注明来源
                - 如果信息不足，明确说明
                - 字数控制在800-1200字
                """, researchTopic, worker1Response, worker2Response);

        String refinerResponse = refinerAgent.prompt()
                .user(refinerPrompt)
                .call()
                .content();

        log.info("✅ Refiner 初稿生成完成(长度: {} 字符)", refinerResponse.length());
        log.info("初稿预览(前500字符):\n{}...", refinerResponse.substring(0, Math.min(500, refinerResponse.length())));

        // ========== 阶段5: Reviewer Agent - 质检初稿 ==========
        log.info("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
        log.info("🔍 阶段5: Reviewer Agent 质量检查");
        log.info("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");

        ChatClient reviewerAgent = ChatClient.builder(chatModel).build();

        String reviewerPrompt = String.format("""
                你是一个研报质量检查员(Reviewer Agent)。

                请审查以下研报初稿，检查:
                1. 是否存在AI幻觉(编造不存在的信息)?
                2. 引用的来源是否清晰准确?
                3. 逻辑是否连贯?
                4. 是否遗漏重要信息?

                初稿内容:
                %s

                原始搜索结果:
                %s

                请用JSON格式返回:
                {
                  "pass": true/false,
                  "score": 0-100,
                  "issues": ["问题1", "问题2", ...],
                  "suggestions": "改进建议"
                }

                如果 pass=false，则需要打回重写。
                """, refinerResponse, worker1Response);

        String reviewerResponse = reviewerAgent.prompt()
                .user(reviewerPrompt)
                .call()
                .content();

        log.info("✅ Reviewer 质检结果:\n{}", reviewerResponse);

        // 简单解析质检结果(实际应用应该用JSON解析)
        boolean reviewPassed = reviewerResponse.toLowerCase().contains("\"pass\": true") ||
                reviewerResponse.toLowerCase().contains("\"pass\":true");

        String finalReport = refinerResponse;

        if (!reviewPassed) {
            log.warn("⚠️ 质检未通过，触发重写流程...");

            // 重写逻辑(简化版: 将reviewer的建议追加到prompt)
            String rewritePrompt = String.format("""
                    请根据以下质检意见，重写研报:

                    原稿:
                    %s

                    质检意见:
                    %s

                    请生成改进后的完整研报(Markdown格式)。
                    """, refinerResponse, reviewerResponse);

            finalReport = refinerAgent.prompt()
                    .user(rewritePrompt)
                    .call()
                    .content();

            log.info("✅ 重写完成");
        } else {
            log.info("✅ 质检通过，无需重写");
        }

        // ========== 阶段6: Finalizer Agent - 保存报告 ==========
        log.info("\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
        log.info("💾 阶段6: Finalizer Agent 保存报告");
        log.info("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");

        // 生成文件名(带时间戳)
        String timestamp = LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyyMMdd_HHmmss"));
        String filename = String.format("%s_研报_%s.md", researchTopic.replace(" ", "_"), timestamp);
        Path reportPath = Paths.get(OUTPUT_DIR, filename);

        try {
            Files.writeString(reportPath, finalReport);
            log.info("✅ 研报保存成功: {}", reportPath.toAbsolutePath());
            log.info("📄 文件大小: {} bytes", Files.size(reportPath));
        } catch (IOException e) {
            log.error("❌ 研报保存失败", e);
        }

        // ========== 总结 ==========
        log.info("\n" +
                "╔═══════════════════════════════════════════════════════════════╗\n" +
                "║                   🎉 工作流执行完成                             ║\n" +
                "╚═══════════════════════════════════════════════════════════════╝");
        log.info("⏰ 结束时间: {}", LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")));
        log.info("📊 最终报告路径: {}", reportPath.toAbsolutePath());

        // 关闭MCP客户端
        try {
            searchClient.closeGracefully();
            filesystemClient.closeGracefully();
            log.info("✅ MCP 客户端关闭成功");
        } catch (Exception e) {
            log.warn("⚠️ MCP 客户端关闭失败: {}", e.getMessage());
        }
    }
}
